// Copyright (c) 2009-2010, Jeremy Brewer
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// * Redistributions of source code must retain the above copyright notice,
// this list of conditions and the following disclaimer.
// * Redistributions in binary form must reproduce the above copyright
// notice, this list of conditions and the following disclaimer in the
// documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.

package galaxie500.astro;

import galaxie500.numeric.Function;
import galaxie500.numeric.Integrator;
import galaxie500.util.MathUtil;
import galaxie500.util.RangeIterator;

/**
 * Class for computing cosmological distances
 * 
 * @author Jeremy Brewer
 */
public class Cosmology {
  private double omegaM;
  private double omegaLambda;
  private double omegaK;
  private double h;

  /** Hubble's constant at z=0 in km/s/MpC */
  private double h0;

  /** Hubble distance in MpC */
  private double hubbleDist;

  /** Speed of light in km/s */
  public static final double SPEED_OF_LIGHT = 2.99792458e5;

  /**
   * Uses latest WMAP 5 cosmology: omega_lambda = 0.72, omega_m = 0.28, and h =
   * 0.71.
   */
  public Cosmology() {
    this(0.28, 0.72, 0.71);
  }

  /**
   * Computes distances for a custom cosmology.
   * 
   * @param omegaM
   * @param omegaLambda
   * @param h
   */
  public Cosmology(double omegaM, double omegaLambda, double h) {
    this.omegaM = omegaM;
    this.omegaLambda = omegaLambda;
    this.omegaK = 1.0 - omegaM - omegaLambda;
    this.h = h;
    this.h0 = h * 100.0;
    this.hubbleDist = SPEED_OF_LIGHT / h0;
  }

  // Accessors.

  public double omegaM() {
    return omegaM;
  }

  public double omegaLambda() {
    return omegaLambda;
  }

  public double omegaK() {
    return omegaK;
  }

  public double h() {
    return h;
  }

  public double h0() {
    return h0;
  }

  public double hubbleDist() {
    return hubbleDist;
  }

  /**
   * Implements the main function to integrate, 1 / E(z).
   */
  private static class InverseE implements Function {
    private Cosmology c;

    public InverseE(Cosmology c) {
      this.c = c;
    }

    public double valueAt(double z) {
      double onePlusZ = 1.0 + z;
      double s = MathUtil.square(onePlusZ);
      double e = c.omegaLambda + s * (c.omegaK + onePlusZ * c.omegaM);
      return 1.0 / Math.sqrt(e);
    }
  }

  /**
   * Implements the co-moving volume element in units Mpc^3 without the solid
   * angle term; you must multiply this by the area of your survey in rad^2.
   */
  private static class ComovingVolElement implements Function {
    private Cosmology c;

    public ComovingVolElement(Cosmology c) {
      this.c = c;
    }

    public double valueAt(double z) {
      double da = c.angularDiameterDist(z);
      return c.hubbleDist * MathUtil.square((1.0 + z) * da) *
          new InverseE(c).valueAt(z);
    }
  }

  /**
   * Wrapper for integrating functions. Primarily abstracted in case we want to
   * apply a global error threshold.
   */
  private double integrate(Function f, double a, double b) {
    Integrator i = new Integrator();
    return i.integrate(f, a, b);
  }

  /**
   * Returns E(z).
   * 
   * @param z Redshift to evaluate at
   * @return The value of E(z)
   */
  public double E(double z) {
    return 1.0 / new InverseE(this).valueAt(z);
  }

  public double comovingLineOfSightDist(double z) {
    double dc = integrate(new InverseE(this), 0.0, z);
    return dc * hubbleDist;
  }

  public double transverseComovingDist(double z) {
    double dc = comovingLineOfSightDist(z);
    if (Math.abs(omegaK) < 1.0e-10) {
      return dc;
    } else if (omegaK > 0.0) {
      double sk = Math.sqrt(omegaK);
      return Math.sinh(sk * dc) / sk;
    } else {
      double sk = Math.sqrt(Math.abs(omegaK));
      return Math.sin(sk * dc) / sk;
    }
  }

  public double angularDiameterDist(double z) {
    double dm = transverseComovingDist(z);
    return dm / (1.0 + z);
  }

  public double luminosityDist(double z) {
    double dm = transverseComovingDist(z);
    return (1.0 + z) * dm;
  }

  // TODO: return value is -Infinity for z = 0.
  public double distanceModulus(double z) {
    final double TEN_PARSECS = 1.0e-5;
    double dl = luminosityDist(z);
    return 5.0 * Math.log10(MathUtil.safeDivide(dl, TEN_PARSECS));
  }

  public double comovingVolElement(double z) {
    return new ComovingVolElement(this).valueAt(z);
  }

  /**
   * Computes the total comoving volume from zmin to zmax with a survey area
   * area in rad^2.
   */
  public double volume(double zmin, double zmax, double area) {
    double vol = integrate(new ComovingVolElement(this), zmin, zmax);
    return vol * area;
  }

  /**
   * Instead of a proper unit test, for now we plot out the standard cosmology
   * from the Hogg paper and inspect it.
   * 
   * @param args Not used
   */
  public static void main(String[] args) {
    Cosmology c = new Cosmology();    
    System.out.printf("# Distances are in units Dh=Hubble dist\n");
    System.out.printf("# z  Dm  Da  Dl  DistModulus  dVc\n");
    
    for (RangeIterator iter = new RangeIterator(0.0, 5.0, 1000);
         !iter.done(); iter.next()) {
      double z = iter.value();
      double dm = c.transverseComovingDist(z) / c.hubbleDist();
      double da = c.angularDiameterDist(z) / c.hubbleDist();
      double dl = c.luminosityDist(z) / c.hubbleDist();
      double distMod = c.distanceModulus(z);
      double dv = c.comovingVolElement(z) / MathUtil.cube(c.hubbleDist());
      System.out.printf("%f  %f  %f  %f  %f  %f\n", z, dm, da, dl,
                        distMod, dv);
    }
  }
}